package com.iris.model;

import org.datavec.api.records.reader.impl.csv.CSVRecordReader;
import org.datavec.api.split.FileSplit;
import org.datavec.api.util.ClassPathResource;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.deeplearning4j.eval.Evaluation;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.SplitTestAndTrain;
import org.nd4j.linalg.dataset.api.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.NormalizerStandardize;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.File;

public class IrisModel {
    /**
     * 加载并预处理训练和测试数据。
     * 从指定文件中读取数据，进行标准化处理，并返回预处理后的数据集。
     *
     * @param filename 数据文件名，位于类路径下。
     * @return 包含预处理后训练和测试数据的 DataSet 对象。
     * @throws Exception 如果在读取文件或处理数据时发生错误。
     */
    public static DataSet getTrainAndTestData(String filename) throws Exception {
        // 初始化 CSV 记录读取器，跳过文件中的指定行数（通常是标题行），并设置分隔符为逗号
        int numLinesToSkip = 1;
        char delimiter = ',';
        CSVRecordReader recordReader = new CSVRecordReader(numLinesToSkip, delimiter);
        recordReader.initialize(new FileSplit(new ClassPathResource(filename).getFile()));

        // 设置标签索引（第5列）和类别数量（3个类别）
        int labelIndex = 4;
        int numClasses = 3;
        int batchSize = 150;

        // 创建数据集迭代器，用于将记录读取器的数据转换为数据集
        DataSetIterator iterator = new RecordReaderDataSetIterator(recordReader, batchSize, labelIndex, numClasses);
        DataSet allData = iterator.next();
        // 打乱数据以确保随机性
        allData.shuffle();

        // 创建标准化器并拟合数据，然后对数据进行标准化处理
        DataNormalization normalizer = new NormalizerStandardize();
        normalizer.fit(allData);
        normalizer.transform(allData);

        // 返回预处理后的数据集
        return allData;
    }

    public static MultiLayerNetwork createModel() {
        MultiLayerConfiguration configuration = new NeuralNetConfiguration
                .Builder()
                .iterations(1000)
                .activation(Activation.TANH)
                .weightInit(WeightInit.XAVIER)
                .learningRate(0.1)
                .regularization(true)
                .l2(0.0001)
                .list()
                .layer(0, new DenseLayer.Builder().nIn(4).nOut(3).build())
                .layer(1, new DenseLayer.Builder().nIn(3).nOut(3).build())
                .layer(2, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                        .activation(Activation.SOFTMAX).nIn(3).nOut(3).build()).backprop(true).pretrain(false).build();
        MultiLayerNetwork model = new MultiLayerNetwork(configuration);
        model.init();
        return model;
    }

    public static void main(String[] args) throws Exception {
        //获取训练和测试数据
        DataSet trainAndTestData = getTrainAndTestData("iris.csv");
        //参数 0.65 表示将数据集按照 65% 和 35% 的比例划分为训练集和测试集。
        SplitTestAndTrain testAndTrain = trainAndTestData.splitTestAndTrain(0.65);
        //训练数据
        DataSet trainingData = testAndTrain.getTrain();
        //测试数据
        DataSet testData = testAndTrain.getTest();
        //创建模型
        MultiLayerNetwork model = createModel();
        //监听器
        model.setListeners(new ScoreIterationListener(100));
        //训练
        model.fit(trainingData);

        //保存训练模型
        // 保存模型
        ModelSerializer.writeModel(model, new File("iris_model.zip"), true);

        //测试训练的结果
        INDArray output = model.output(testData.getFeatureMatrix());
        Evaluation eval = new Evaluation(3);
        eval.eval(testData.getLabels(), output);
        System.out.println(eval.stats());
    }
}
